package com.example.terminalservice.utils;

import com.example.terminalservice.pojo.dto.ConnInfoDTO;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jcraft.jsch.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;

import javax.websocket.*;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * @author 28142
 * @description TerminalWebSocketHandler
 * @date 2023/8/7 23:52
 */
@Component
@ServerEndpoint("/terminal/{username}")
public class TerminalWebSocketHandler {
    private ObjectMapper objectMapper = new ObjectMapper();
    private static final ConcurrentMap<String, ChannelShell> shellChannels = new ConcurrentHashMap<>();

    @OnOpen
    public void onOpen(Session session, @PathParam("username") String username) {
        System.out.println("------------websocket连接成功------------------");
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        // 获取url中的username
        String username = session.getPathParameters().get("username");
        System.out.println("===============接受消息=================");
        System.out.println(username + ": " + message);
        try {
            ConnInfoDTO connInfoDTO = objectMapper.readValue(message, ConnInfoDTO.class);
            System.out.println("连接对象");
            System.out.println(connInfoDTO);
            handleConnInfo(connInfoDTO, session, username);
        } catch (JsonProcessingException e) {
            System.out.println("消息命令");
            handleTextMessage(message, session, username);
        }
    }

    private void handleConnInfo(ConnInfoDTO connInfoDTO, Session session, String username) {
        // 处理收到的连接信息
        System.out.println("==========处理连接信息==========");
        System.out.println(connInfoDTO);
        try {
            JSch jsch = new JSch();
            com.jcraft.jsch.Session sshSession = jsch.getSession(connInfoDTO.getUsername(), connInfoDTO.getHost(), connInfoDTO.getPort());
            sshSession.setPassword(connInfoDTO.getPassword());
            sshSession.setConfig("StrictHostKeyChecking", "no");
            sshSession.connect();
            System.out.println("***********建立ssh连接***********");
            System.out.println(sshSession);

            // 创建会话
            ChannelShell channel = (ChannelShell) sshSession.openChannel("shell");
            channel.connect();

            // 获取输入输出流
            InputStream inputStream = channel.getInputStream();
            OutputStream outputStream = channel.getOutputStream();
            // 在onOpen方法中添加一个线程用于读取输出流
            Thread readThread = new Thread(() -> {
                try {
                    //循环读取
                    byte[] buffer = new byte[1024];
                    int i = 0;
                    //如果没有数据来，线程会一直阻塞在这个地方等待数据。
                    while ((i = inputStream.read(buffer)) != -1) {
                        session.getBasicRemote().sendText(new TextMessage(Arrays.copyOfRange(buffer, 0, i)).getPayload());
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    sshSession.disconnect();
                    channel.disconnect();
                    try {
                        inputStream.close();
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                }
            });
            readThread.start();
            shellChannels.put(username, channel);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private void handleTextMessage(String message, Session session, String username) {
        // 处理收到的普通文本消息
        // 根据username从userSessions获取到对应的会话
        ChannelShell channel = shellChannels.get(username);
        try {
            if (session.isOpen()) {
                OutputStream outputStream = channel.getOutputStream();
                outputStream.write(message.getBytes());
                outputStream.write('\n');
                outputStream.flush();
            }
        } catch (IOException exception) {
            throw new RuntimeException(exception);
        }
    }

    @OnClose
    public void onClose(Session session) {
        System.out.println("*******************close*******************");
        // 获取url中的username
        String username = session.getPathParameters().get("username");
        // 根据username从shellChannels获取对应的连接
        ChannelShell channel = shellChannels.get(username);


        try {
            if (channel != null && channel.isConnected()) {
                channel.disconnect();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        error.printStackTrace();
    }
}
